# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import torch
from torch import nn
from feature_extract.model.fast_reid.resnet import *
from feature_extract.model.fast_reid.layers.pooling import GeneralizedMeanPooling
from feature_extract.registry import MODEL_REGISTRY
from feature_extract.model.fast_reid.head.embedding_head import *


@MODEL_REGISTRY.register('fast_reid_resnet')
class Baseline(nn.Module):
    """
    Baseline architecture. Any models that contains the following two components:
    1. Per-image feature extraction (aka backbone)
    2. Per-image feature aggregation and loss computation
    """

    def __init__(
            self,
            with_ibn,
            with_nl,
            depth,
            pool_type,
            feat_dim=2048,
            embedding_dim=0,
            num_classes=0,
            neck_feat="before",
            cls_type="Linear",
            scale=1,
            margin=0,
            with_bnneck=False,
            norm_type="BN",
            pixel_std=None,
            loss_kwargs=None,
            pixel_mean=None,
            **kwargs
    ):
        """
        NOTE: this interface is experimental.

        Args:
            backbone:
            heads:
            pixel_mean:
            pixel_std:
        """
        super().__init__()
        # backbone
        if pixel_std is None:
            pixel_std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
        if pixel_mean is None:
            pixel_mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
        self.backbone = build_resnet_backbone(with_ibn=with_ibn, with_nl=with_nl, depth=depth)
        # head
        self.heads = EmbeddingHead(feat_dim,
                                   embedding_dim,
                                   num_classes,
                                   neck_feat,
                                   pool_type,
                                   cls_type,
                                   scale,
                                   margin,
                                   with_bnneck,
                                   norm_type)

        self.loss_kwargs = loss_kwargs

        self.register_buffer('pixel_mean', torch.Tensor(pixel_mean).view(1, -1, 1, 1), False)
        self.register_buffer('pixel_std', torch.Tensor(pixel_std).view(1, -1, 1, 1), False)

    def device(self):
        return self.pixel_mean.device

    def forward(self, batched_inputs):
        features = self.backbone(batched_inputs)
        outputs = self.heads(features)
        return outputs
